{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "163ff75c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _components.producer_decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "563b0aba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "import asyncio\n",
    "import functools\n",
    "import logging\n",
    "import random\n",
    "import time\n",
    "from asyncio import iscoroutinefunction  # do not use the version from inspect\n",
    "from dataclasses import dataclass\n",
    "from functools import partial\n",
    "from inspect import Parameter\n",
    "from typing import *\n",
    "\n",
    "from aiokafka import AIOKafkaProducer\n",
    "from aiokafka.errors import KafkaTimeoutError, RequestTimedOutError\n",
    "from aiokafka.producer.message_accumulator import BatchBuilder\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from fastkafka._components.logger import get_logger, cached_log\n",
    "from fastkafka._components.meta import export\n",
    "from fastkafka._components.helpers import remove_suffix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e222a73f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import unittest\n",
    "from contextlib import asynccontextmanager, contextmanager\n",
    "from itertools import product\n",
    "from unittest.mock import ANY, Mock, call\n",
    "from _pytest import monkeypatch\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from fastkafka._components.logger import suppress_timestamps\n",
    "from fastkafka.encoder import avro_encoder, json_encoder\n",
    "from fastkafka._testing.in_memory_broker import InMemoryBroker, InMemoryProducer, InMemoryConsumer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "436c6e0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "logger = get_logger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fae4d87e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: ok\n"
     ]
    }
   ],
   "source": [
    "suppress_timestamps()\n",
    "logger = get_logger(__name__, level=20)\n",
    "logger.info(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b4e5e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "BaseSubmodel = TypeVar(\"BaseSubmodel\", bound=Union[List[BaseModel], BaseModel])\n",
    "BaseSubmodel\n",
    "\n",
    "\n",
    "@dataclass\n",
    "@export(\"fastkafka\")\n",
    "class KafkaEvent(Generic[BaseSubmodel]):\n",
    "    \"\"\"\n",
    "    A generic class for representing Kafka events. Based on BaseSubmodel, bound to pydantic.BaseModel\n",
    "\n",
    "    Attributes:\n",
    "        message (BaseSubmodel): The message contained in the Kafka event, can be of type pydantic.BaseModel.\n",
    "        key (bytes, optional): The optional key used to identify the Kafka event.\n",
    "    \"\"\"\n",
    "\n",
    "    message: BaseSubmodel\n",
    "    key: Optional[bytes] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e0166e",
   "metadata": {},
   "outputs": [],
   "source": [
    "event = KafkaEvent(\"Some message\")\n",
    "assert event.message == \"Some message\"\n",
    "assert event.key == None\n",
    "\n",
    "event = KafkaEvent(\"Some message\", b\"123\")\n",
    "assert event.message == \"Some message\"\n",
    "assert event.key == b\"123\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e7de730",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def unwrap_from_kafka_event(var_type: Union[Type, Parameter]) -> Union[Type, Parameter]:\n",
    "    \"\"\"\n",
    "    Unwraps the type from a KafkaEvent.\n",
    "\n",
    "    Args:\n",
    "        var_type: Type to unwrap.\n",
    "\n",
    "    Returns:\n",
    "        Type: Unwrapped type if the given type is a KafkaEvent, otherwise returns the same type.\n",
    "\n",
    "    Example:\n",
    "        - Input: KafkaEvent[str]\n",
    "          Output: str\n",
    "        - Input: int\n",
    "          Output: int\n",
    "    \"\"\"\n",
    "    if hasattr(var_type, \"__origin__\") and var_type.__origin__ == KafkaEvent:\n",
    "        return var_type.__args__[0]  # type: ignore\n",
    "    else:\n",
    "        return var_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7277bcf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert unwrap_from_kafka_event(KafkaEvent[int]) == int\n",
    "assert unwrap_from_kafka_event(int) == int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d981cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "ProduceReturnTypes = Union[\n",
    "    BaseModel, KafkaEvent[BaseModel], List[BaseModel], KafkaEvent[List[BaseModel]]\n",
    "]\n",
    "\n",
    "ProduceCallable = Union[\n",
    "    Callable[..., ProduceReturnTypes], Callable[..., Awaitable[ProduceReturnTypes]]\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a66cc7a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # | export\n",
    "\n",
    "\n",
    "# def _to_json_utf8(o: Any) -> bytes:\n",
    "#     \"\"\"Converts to JSON and then encodes with UTF-8\"\"\"\n",
    "#     if hasattr(o, \"json\"):\n",
    "#         return o.json().encode(\"utf-8\")  # type: ignore\n",
    "#     else:\n",
    "#         return json.dumps(o).encode(\"utf-8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bd6af9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# assert _to_json_utf8({\"a\": 1, \"b\": [2, 3]}) == b'{\"a\": 1, \"b\": [2, 3]}'\n",
    "\n",
    "\n",
    "class ExampleMsg(BaseModel):\n",
    "    name: str = Field()\n",
    "    age: int\n",
    "\n",
    "\n",
    "# assert _to_json_utf8(ExampleMsg(name=\"Davor\", age=12)) == b'{\"name\": \"Davor\", \"age\": 12}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eb9ce04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _wrap_in_event(\n",
    "    message: Union[BaseModel, List[BaseModel], KafkaEvent]\n",
    ") -> KafkaEvent:\n",
    "    return message if type(message) == KafkaEvent else KafkaEvent(message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50cab521",
   "metadata": {},
   "outputs": [],
   "source": [
    "message = ExampleMsg(name=\"Davor\", age=12)\n",
    "wrapped = _wrap_in_event(message)\n",
    "\n",
    "assert type(wrapped) == KafkaEvent\n",
    "assert wrapped.message == message\n",
    "assert wrapped.key == None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d82893da",
   "metadata": {},
   "outputs": [],
   "source": [
    "message = KafkaEvent(ExampleMsg(name=\"Davor\", age=12), b\"123\")\n",
    "wrapped = _wrap_in_event(message)\n",
    "\n",
    "assert type(wrapped) == KafkaEvent\n",
    "assert wrapped.message == message.message\n",
    "assert wrapped.key == b\"123\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18abd215",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def release_callback(\n",
    "    fut: asyncio.Future, topic: str, wrapped_val: KafkaEvent[BaseModel]\n",
    ") -> None:\n",
    "    if fut.exception() is not None:\n",
    "        cached_log(\n",
    "            logger,\n",
    "            f\"release_callback(): Exception {fut.exception()=}, raised when producing {wrapped_val.message=} to {topic=}\",\n",
    "            level=logging.WARNING,\n",
    "            timeout=1,\n",
    "            log_id=\"release_callback()\"\n",
    "        )\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd83badb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "async def produce_single(  # type: ignore\n",
    "    producer: AIOKafkaProducer,\n",
    "    topic: str,\n",
    "    encoder_fn: Callable[[BaseModel], bytes],\n",
    "    wrapped_val: KafkaEvent[BaseModel],\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Sends a single message to the Kafka producer.\n",
    "\n",
    "    Args:\n",
    "        producer (AIOKafkaProducer): The Kafka producer object.\n",
    "        topic (str): The topic to which the message will be sent.\n",
    "        encoder_fn (Callable[[BaseModel], bytes]): The encoding function to encode the message.\n",
    "        wrapped_val (KafkaEvent[BaseModel]): The wrapped Kafka event containing the message.\n",
    "    \"\"\"\n",
    "    while True:\n",
    "        try:\n",
    "            fut = await producer.send(\n",
    "                topic, encoder_fn(wrapped_val.message), key=wrapped_val.key\n",
    "            )\n",
    "            fut.add_done_callback(partial(release_callback, topic=topic, wrapped_val=wrapped_val))\n",
    "            break\n",
    "        except KafkaTimeoutError as e:\n",
    "            logger.warning(f\"produce_single(): Exception {e=} raised when producing {wrapped_val.message} to {topic=}, sleeping for 1 second and retrying..\")\n",
    "            await asyncio.sleep(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7c363dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FakeProducer:\n",
    "    def __init__(self, return_future: asyncio.Future = None):\n",
    "        self.counter = 0\n",
    "        if return_future is None:\n",
    "            return_future = asyncio.Future()\n",
    "            return_future.set_result(\"Some result\")\n",
    "\n",
    "        self.return_future = return_future\n",
    "\n",
    "    async def send(self, *args, **kwargs):\n",
    "        if self.counter < 5:\n",
    "            self.counter += 1\n",
    "            raise KafkaTimeoutError()\n",
    "        else:\n",
    "            return self.return_future\n",
    "\n",
    "    async def send_batch(self, *args, **kwargs):\n",
    "        return await self.send()\n",
    "\n",
    "    def create_batch(self):\n",
    "        return unittest.mock.MagicMock()\n",
    "\n",
    "    def add_done_callback(self, *args):\n",
    "        return\n",
    "\n",
    "    async def partitions_for(self, *args):\n",
    "        return [\"partition_1\", \"partition_2\"]\n",
    "    \n",
    "    def start(*args, **kwargs):\n",
    "        return\n",
    "    \n",
    "    def stop(*args, **kwargs):\n",
    "        return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3fb499d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n"
     ]
    }
   ],
   "source": [
    "await produce_single(\n",
    "    FakeProducer(),\n",
    "    topic=\"test_topic\",\n",
    "    encoder_fn=json_encoder,\n",
    "    wrapped_val=KafkaEvent(message=ExampleMsg(name=\"Davor\", age=12), key=b\"test\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccb950e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: produce_single(): Exception e=KafkaTimeoutError() raised when producing name='Davor' age=12 to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: release_callback(): Exception fut.exception()=RequestTimedOutError(), raised when producing wrapped_val.message=ExampleMsg(name='Davor', age=12) to topic='test_topic'\n"
     ]
    }
   ],
   "source": [
    "timeout_future = asyncio.Future()\n",
    "timeout_future.set_exception(RequestTimedOutError())\n",
    "\n",
    "await produce_single(\n",
    "    FakeProducer(return_future=timeout_future),\n",
    "    topic=\"test_topic\",\n",
    "    encoder_fn=json_encoder,\n",
    "    wrapped_val=KafkaEvent(message=ExampleMsg(name=\"Davor\", age=12), key=b\"test\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6a9221e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "with InMemoryBroker() as broker:\n",
    "    ProducerClass = InMemoryProducer(broker)\n",
    "    producer = ProducerClass()\n",
    "    await producer.start()\n",
    "\n",
    "    await produce_single(\n",
    "        producer,\n",
    "        topic=\"test_topic\",\n",
    "        encoder_fn=json_encoder,\n",
    "        wrapped_val=KafkaEvent(message=ExampleMsg(name=\"Davor\", age=12), key=b\"test\"),\n",
    "    )\n",
    "\n",
    "    await producer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e20ee70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "async def send_batch(  # type: ignore\n",
    "    producer: AIOKafkaProducer, topic: str, batch: BatchBuilder, key: Optional[bytes]\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Sends a batch of messages to the Kafka producer.\n",
    "\n",
    "    Args:\n",
    "        producer (AIOKafkaProducer): The Kafka producer object.\n",
    "        topic (str): The topic to which the messages will be sent.\n",
    "        batch (BatchBuilder): The batch builder object containing the messages.\n",
    "        key (Optional[bytes]): The optional key used to identify the batch of messages.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    partitions = await producer.partitions_for(topic)\n",
    "    if key == None:\n",
    "        partition = random.choice(tuple(partitions))  # nosec\n",
    "    else:\n",
    "        partition = producer._partition(topic, None, None, None, key, None)\n",
    "    while True:\n",
    "        try:\n",
    "            await producer.send_batch(batch, topic, partition=partition)\n",
    "            break\n",
    "        except KafkaTimeoutError as e:\n",
    "            logger.warning(f\"send_batch(): Exception {e} raised when producing {batch} to {topic=}, sleeping for 1 second and retrying..\")\n",
    "            await asyncio.sleep(1)\n",
    "    \n",
    "\n",
    "\n",
    "async def produce_batch(  # type: ignore\n",
    "    producer: AIOKafkaProducer,\n",
    "    topic: str,\n",
    "    encoder_fn: Callable[[BaseModel], bytes],\n",
    "    wrapped_val: KafkaEvent[List[BaseModel]],\n",
    ") -> ProduceReturnTypes:\n",
    "    \"\"\"\n",
    "    Sends a batch of messages to the Kafka producer.\n",
    "\n",
    "    Args:\n",
    "        producer (AIOKafkaProducer): The Kafka producer object.\n",
    "        topic (str): The topic to which the messages will be sent.\n",
    "        encoder_fn (Callable[[BaseModel], bytes]): The encoding function to encode the messages.\n",
    "        wrapped_val (KafkaEvent[List[BaseModel]]): The wrapped Kafka event containing the list of messages.\n",
    "\n",
    "    Returns:\n",
    "        ProduceReturnTypes: The return value from the decorated function.\n",
    "    \"\"\"\n",
    "    batch = producer.create_batch()\n",
    "\n",
    "    for message in wrapped_val.message:\n",
    "        metadata = batch.append(\n",
    "            key=wrapped_val.key,\n",
    "            value=encoder_fn(message),\n",
    "            timestamp=int(time.time() * 1000),\n",
    "        )\n",
    "        if metadata == None:\n",
    "            # send batch\n",
    "            await send_batch(producer, topic, batch, wrapped_val.key)\n",
    "            # create new batch\n",
    "            batch = producer.create_batch()\n",
    "            batch.append(\n",
    "                key=None, value=encoder_fn(message), timestamp=int(time.time() * 1000)\n",
    "            )\n",
    "\n",
    "    await send_batch(producer, topic, batch, wrapped_val.key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7725fd2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[WARNING] __main__: send_batch(): Exception KafkaTimeoutError raised when producing <MagicMock id='140539229227984'> to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: send_batch(): Exception KafkaTimeoutError raised when producing <MagicMock id='140539229227984'> to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: send_batch(): Exception KafkaTimeoutError raised when producing <MagicMock id='140539229227984'> to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: send_batch(): Exception KafkaTimeoutError raised when producing <MagicMock id='140539229227984'> to topic='test_topic', sleeping for 1 second and retrying..\n",
      "[WARNING] __main__: send_batch(): Exception KafkaTimeoutError raised when producing <MagicMock id='140539229227984'> to topic='test_topic', sleeping for 1 second and retrying..\n"
     ]
    }
   ],
   "source": [
    "msgs = [ExampleMsg(name=\"Davor\", age=12) for _ in range(500)]\n",
    "\n",
    "await produce_batch(\n",
    "    FakeProducer(),\n",
    "    topic=\"test_topic\",\n",
    "    encoder_fn=json_encoder,\n",
    "    wrapped_val=KafkaEvent(message=msgs, key=None),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee8b2b92",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "msgs = [ExampleMsg(name=\"Davor\", age=12) for _ in range(500)]\n",
    "    \n",
    "with InMemoryBroker() as broker:\n",
    "    ProducerClass = InMemoryProducer(broker)\n",
    "    producer = ProducerClass()\n",
    "    await producer.start()\n",
    "\n",
    "    await produce_batch(\n",
    "        producer,\n",
    "        topic=\"test_topic\",\n",
    "        encoder_fn=json_encoder,\n",
    "        wrapped_val=KafkaEvent(message=msgs, key=b\"test\"),\n",
    "    )\n",
    "\n",
    "    await producer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87f60d23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def producer_decorator(\n",
    "    producer_store: Dict[str, Any],\n",
    "    func: ProduceCallable,\n",
    "    topic_key: str,\n",
    "    encoder_fn: Callable[[BaseModel], bytes],\n",
    ") -> ProduceCallable:\n",
    "    \"\"\"\n",
    "    Decorator for Kafka producer functions.\n",
    "\n",
    "    Args:\n",
    "        producer_store (Dict[str, Any]): Dictionary to store the Kafka producer objects.\n",
    "        func (ProduceCallable): The function to be decorated.\n",
    "        topic_key (str): The key used to identify the topic.\n",
    "        encoder_fn (Callable[[BaseModel], bytes]): The encoding function to encode the messages.\n",
    "\n",
    "    Returns:\n",
    "        ProduceCallable: The decorated function.\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If the decorated function is synchronous.\n",
    "    \"\"\"\n",
    "\n",
    "    @functools.wraps(func)\n",
    "    async def _produce_async(\n",
    "        *args: List[Any],\n",
    "        topic_key: str = topic_key,\n",
    "        encoder_fn: Callable[[BaseModel], bytes] = encoder_fn,\n",
    "        producer_store: Dict[str, Any] = producer_store,\n",
    "        f: Callable[..., Awaitable[ProduceReturnTypes]] = func,  # type: ignore\n",
    "        **kwargs: Any,\n",
    "    ) -> ProduceReturnTypes:\n",
    "        return_val = await f(*args, **kwargs)\n",
    "        wrapped_val = _wrap_in_event(return_val)\n",
    "        _, producer, _, _ = producer_store[topic_key]\n",
    "        topic = remove_suffix(topic_key)\n",
    "\n",
    "        if isinstance(wrapped_val.message, list):\n",
    "            await produce_batch(producer, topic, encoder_fn, wrapped_val)\n",
    "        else:\n",
    "            await produce_single(producer, topic, encoder_fn, wrapped_val)\n",
    "        return return_val\n",
    "\n",
    "    if not iscoroutinefunction(func):\n",
    "        raise ValueError(\n",
    "            \"Synchronous functions are not supported for produce operation\"\n",
    "        )\n",
    "\n",
    "    return _produce_async"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76940b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MockMsg(BaseModel):\n",
    "    name: str = \"Micky Mouse\"\n",
    "    id: int = 123\n",
    "\n",
    "\n",
    "mock_msg = MockMsg()\n",
    "\n",
    "topic = \"test_topic_1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d950429",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def _f() -> None:\n",
    "#     print(\"Mock called\")\n",
    "    loop = asyncio.get_running_loop()\n",
    "\n",
    "    # Create a new Future object.\n",
    "    return loop.create_future()\n",
    "\n",
    "\n",
    "@contextmanager\n",
    "def mock_InMemoryProducer_send() -> Generator[Mock, None, None]:\n",
    "    \"\"\"Mocks **send** method of **InMemoryProducer**\"\"\"\n",
    "    with unittest.mock.patch(\n",
    "        \"fastkafka._testing.in_memory_broker.InMemoryProducer.send\"\n",
    "    ) as mock:\n",
    "        mock.return_value = asyncio.create_task(_f())\n",
    "\n",
    "        yield mock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d66e6b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@asynccontextmanager\n",
    "async def mock_producer_send_env() -> AsyncGenerator[\n",
    "    Tuple[Mock, AIOKafkaProducer], None\n",
    "]:\n",
    "    try:\n",
    "        with mock_InMemoryProducer_send() as send_mock:\n",
    "            with InMemoryBroker() as broker:\n",
    "                ProducerClass = InMemoryProducer(broker)\n",
    "                producer = ProducerClass()\n",
    "                await producer.start()\n",
    "\n",
    "                yield send_mock, producer\n",
    "    finally:\n",
    "        await producer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60e06f18",
   "metadata": {},
   "outputs": [],
   "source": [
    "@asynccontextmanager\n",
    "async def mock_producer_batch_env() -> AsyncGenerator[\n",
    "    Tuple[Mock, AIOKafkaProducer], None\n",
    "]:\n",
    "    try:\n",
    "        with unittest.mock.patch(\n",
    "            \"fastkafka._testing.in_memory_broker.InMemoryProducer.send_batch\"\n",
    "        ) as send_batch_mock, unittest.mock.patch(\n",
    "            \"fastkafka._testing.in_memory_broker.InMemoryProducer.create_batch\"\n",
    "        ) as create_batch_mock:\n",
    "            batch_mock = Mock()\n",
    "            create_batch_mock.return_value = batch_mock\n",
    "            send_batch_mock.return_value = asyncio.create_task(_f())\n",
    "            with InMemoryBroker() as broker:\n",
    "                ProducerClass = InMemoryProducer(broker)\n",
    "                producer = ProducerClass()\n",
    "                await producer.start()\n",
    "\n",
    "                yield batch_mock, send_batch_mock, producer\n",
    "    finally:\n",
    "        await producer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4384ad73",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing with: is_sync=False , encoder_fn=<function json_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "Testing with: is_sync=False , encoder_fn=<function avro_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n"
     ]
    }
   ],
   "source": [
    "async def func_async(mock_msg: MockMsg) -> MockMsg:\n",
    "    return mock_msg\n",
    "\n",
    "\n",
    "def func_sync(mock_msg: MockMsg) -> MockMsg:\n",
    "    return mock_msg\n",
    "\n",
    "\n",
    "is_sync = False\n",
    "for encoder_fn in [json_encoder, avro_encoder]:\n",
    "    print(f\"Testing with: {is_sync=} , {encoder_fn=}\")\n",
    "    async with mock_producer_send_env() as (send_mock, producer):\n",
    "        test_func = producer_decorator(\n",
    "            {topic: (None, producer, None, None)},\n",
    "            func_sync if is_sync else func_async,\n",
    "            topic,\n",
    "            encoder_fn=encoder_fn,\n",
    "        )\n",
    "\n",
    "        assert iscoroutinefunction(test_func) != is_sync\n",
    "\n",
    "        value = test_func(mock_msg) if is_sync else await test_func(mock_msg)\n",
    "\n",
    "        send_mock.assert_called_once_with(remove_suffix(topic), encoder_fn(mock_msg), key=None)\n",
    "\n",
    "        assert value == mock_msg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bec5c19",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing with: is_sync=False , encoder_fn=<function json_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "Testing with: is_sync=False , encoder_fn=<function avro_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n"
     ]
    }
   ],
   "source": [
    "test_key = b\"key\"\n",
    "\n",
    "\n",
    "async def func_async(mock_msg: MockMsg) -> KafkaEvent[MockMsg]:\n",
    "    return KafkaEvent(mock_msg, test_key)\n",
    "\n",
    "\n",
    "def func_sync(mock_msg: MockMsg) -> KafkaEvent[MockMsg]:\n",
    "    return KafkaEvent(mock_msg, test_key)\n",
    "\n",
    "\n",
    "is_sync = False\n",
    "for encoder_fn in [json_encoder, avro_encoder]:\n",
    "    print(f\"Testing with: {is_sync=} , {encoder_fn=}\")\n",
    "    async with mock_producer_send_env() as (send_mock, producer):\n",
    "        test_func = producer_decorator(\n",
    "            {topic: (None, producer, None, None)},\n",
    "            func_sync if is_sync else func_async,\n",
    "            topic,\n",
    "            encoder_fn=encoder_fn,\n",
    "        )\n",
    "\n",
    "        assert iscoroutinefunction(test_func) != is_sync\n",
    "\n",
    "        value = test_func(mock_msg) if is_sync else await test_func(mock_msg)\n",
    "\n",
    "        send_mock.assert_called_once_with(remove_suffix(topic), encoder_fn(mock_msg), key=test_key)\n",
    "\n",
    "        assert value == KafkaEvent(mock_msg, test_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8df14a18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing with: is_sync=False , encoder_fn=<function json_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "Testing with: is_sync=False , encoder_fn=<function avro_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n"
     ]
    }
   ],
   "source": [
    "batch_size = 123\n",
    "\n",
    "\n",
    "async def func_async(mock_msg: MockMsg) -> List[MockMsg]:\n",
    "    return [mock_msg] * batch_size\n",
    "\n",
    "\n",
    "def func_sync(mock_msg: MockMsg) -> List[MockMsg]:\n",
    "    return [mock_msg] * batch_size\n",
    "\n",
    "\n",
    "is_sync = False\n",
    "for encoder_fn in [json_encoder, avro_encoder]:\n",
    "    print(f\"Testing with: {is_sync=} , {encoder_fn=}\")\n",
    "    async with mock_producer_batch_env() as (\n",
    "        batch_mock,\n",
    "        send_batch_mock,\n",
    "        producer,\n",
    "    ):\n",
    "        test_func = producer_decorator(\n",
    "            {topic: (None, producer, None, None)},\n",
    "            func_sync if is_sync else func_async,\n",
    "            topic,\n",
    "            encoder_fn=encoder_fn,\n",
    "        )\n",
    "\n",
    "        assert iscoroutinefunction(test_func) != is_sync\n",
    "\n",
    "        value = test_func(mock_msg) if is_sync else await test_func(mock_msg)\n",
    "\n",
    "        batch_mock.append.assert_has_calls(\n",
    "            [call(key=None, value=encoder_fn(mock_msg), timestamp=ANY)] * batch_size\n",
    "        )\n",
    "        send_batch_mock.assert_called_once_with(batch_mock, remove_suffix(topic), partition=0)\n",
    "\n",
    "        assert value == [mock_msg] * batch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beb344fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing with: is_sync=False , encoder_fn=<function json_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "Testing with: is_sync=False , encoder_fn=<function avro_encoder>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n"
     ]
    }
   ],
   "source": [
    "batch_size = 123\n",
    "test_key = b\"key\"\n",
    "\n",
    "\n",
    "async def func_async(mock_msg: MockMsg) -> KafkaEvent[List[MockMsg]]:\n",
    "    return KafkaEvent([mock_msg] * batch_size, test_key)\n",
    "\n",
    "\n",
    "def func_sync(mock_msg: MockMsg) -> KafkaEvent[List[MockMsg]]:\n",
    "    return KafkaEvent([mock_msg] * batch_size, test_key)\n",
    "\n",
    "\n",
    "is_sync = False\n",
    "for encoder_fn in [json_encoder, avro_encoder]:\n",
    "    print(f\"Testing with: {is_sync=} , {encoder_fn=}\")\n",
    "    async with mock_producer_batch_env() as (batch_mock, send_batch_mock, producer):\n",
    "        test_func = producer_decorator(\n",
    "            {topic: (None, producer, None, None)},\n",
    "            func_sync if is_sync else func_async,\n",
    "            topic,\n",
    "            encoder_fn=encoder_fn,\n",
    "        )\n",
    "\n",
    "        assert iscoroutinefunction(test_func) != is_sync\n",
    "\n",
    "        value = test_func(mock_msg) if is_sync else await test_func(mock_msg)\n",
    "\n",
    "        batch_mock.append.assert_has_calls(\n",
    "            [call(key=test_key, value=encoder_fn(mock_msg), timestamp=ANY)] * batch_size\n",
    "        )\n",
    "\n",
    "        send_batch_mock.assert_called_once_with(batch_mock, remove_suffix(topic), partition=0)\n",
    "\n",
    "        assert value == KafkaEvent([mock_msg] * batch_size, test_key)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
